import numpy as np
from get_data.origin_data import *
from indicator import indicator_merge

# 行列对齐(打印调试用)
pd.set_option('display.unicode.ambiguous_as_wide', True)
pd.set_option('display.unicode.east_asian_width', True)
pd.set_option('display.width', None)
# pd.set_option('display.max_rows', None)
pd.set_option('display.max_columns', None)

# 获取仓位情况
def combine_position():
    df_RSJ_5D = pd.read_excel('因子回测结果/RSJ_5D回测.xlsx', sheet_name='Roll')[['交易日期', '指数净值', '持仓状态']]\
        .rename(columns={'持仓状态':'RSJ_5D持仓状态'})
    df_cov_positive_ret_pct_5D = pd.read_excel('因子回测结果/cov_+ret_pct_5D回测.xlsx', sheet_name='Roll')[['交易日期', '持仓状态']]\
        .rename(columns={'持仓状态':'cov_+ret_pct_5D持仓状态'})
    df_covprice_mid = pd.read_excel('因子回测结果/covprice_mid回测.xlsx', sheet_name='MACD')[['交易日期', '持仓状态']]\
        .rename(columns={'持仓状态':'covprice_mid持仓状态'})
    df_IVdelta_mid_5D = pd.read_excel('因子回测结果/IVdelta_mid_5D回测.xlsx', sheet_name='Roll')[['交易日期', '持仓状态']]\
        .rename(columns={'持仓状态':'IVdelta_mid_5D持仓状态'})
    df_YTMdelta_mid = pd.read_excel('因子回测结果/YTMdelta_mid回测.xlsx', sheet_name='MACD')[['交易日期', '持仓状态']]\
        .rename(columns={'持仓状态':'YTMdelta_mid持仓状态'})
    df_IVnan_pct = pd.read_excel('因子回测结果/IVnan_pct回测.xlsx', sheet_name='MACD')[['交易日期', '持仓状态']]\
        .rename(columns={'持仓状态':'IVnan_pct持仓状态'})

    df = pd.merge(df_RSJ_5D, df_cov_positive_ret_pct_5D, how = 'inner', on = '交易日期')
    df = pd.merge(df, df_covprice_mid, how = 'inner', on = '交易日期')
    df = pd.merge(df, df_IVdelta_mid_5D, how='inner', on='交易日期')
    df = pd.merge(df, df_YTMdelta_mid, how='inner', on='交易日期')
    df = pd.merge(df, df_IVnan_pct, how='inner', on='交易日期')

    df['总持仓数'] = df['RSJ_5D持仓状态'] + df['cov_+ret_pct_5D持仓状态'] + df['covprice_mid持仓状态'] \
                + df['IVdelta_mid_5D持仓状态'] + df['YTMdelta_mid持仓状态'] + df['IVnan_pct持仓状态']

    df['等权仓位占比'] = df['总持仓数'] / 6
    df['分权仓位占比'] = df['IVdelta_mid_5D持仓状态'] * 0.25 + (df['RSJ_5D持仓状态'] + df['cov_+ret_pct_5D持仓状态'] + \
                        df['covprice_mid持仓状态'] +df['YTMdelta_mid持仓状态'] + df['IVnan_pct持仓状态']) * 0.15

    return df

# 组合择时设定1
def combine_time_decision_1():
    df = combine_position()[['交易日期', '指数净值', 'RSJ_5D持仓状态', 'cov_+ret_pct_5D持仓状态', 'covprice_mid持仓状态',
                             'IVdelta_mid_5D持仓状态', 'YTMdelta_mid持仓状态', 'IVnan_pct持仓状态', '总持仓数']]
    df['组合持仓状态'] = (df['总持仓数'] >= 4).astype(int)
    df['日收益率'] = 0
    for i in range(1, len(df)):
        if df.loc[i, '组合持仓状态'] == 1:
            df.loc[i, '日收益率'] = df.loc[i, '指数净值'] / df.loc[i - 1, '指数净值'] - 1
    df['组合累积净值'] = np.nan
    df.loc[0, '组合累积净值'] = 1
    for i in range(1, len(df)):
        df.loc[i, '组合累积净值'] = df.loc[i - 1, '组合累积净值'] * (1 + df.loc[i, '日收益率'])
    return df

# 组合择时设定2
def combine_time_decision_2():
    df = combine_position()[['交易日期', '指数净值', 'RSJ_5D持仓状态', 'cov_+ret_pct_5D持仓状态', 'covprice_mid持仓状态',
                             'IVdelta_mid_5D持仓状态', 'YTMdelta_mid持仓状态', 'IVnan_pct持仓状态', '总持仓数']]
    df.loc[0, '组合持仓状态'] = (df.loc[0, '总持仓数'] >= 4).astype(int)
    for i in range(1, len(df)):
        if df.loc[i, '总持仓数'] >= 4:
            df.loc[i, '组合持仓状态'] = 1
        elif df.loc[i, '总持仓数'] <= 2:
            df.loc[i, '组合持仓状态'] = 0
        else:
            df.loc[i, '组合持仓状态'] = df.loc[i - 1, '组合持仓状态']

    df['日收益率'] = 0
    for i in range(1, len(df)):
        if df.loc[i, '组合持仓状态'] == 1:
            df.loc[i, '日收益率'] = df.loc[i, '指数净值'] / df.loc[i - 1, '指数净值'] - 1
    df['组合累积净值'] = np.nan
    df.loc[0, '组合累积净值'] = 1
    for i in range(1, len(df)):
        df.loc[i, '组合累积净值'] = df.loc[i - 1, '组合累积净值'] * (1 + df.loc[i, '日收益率'])
    return df

# 组合择时设定3
def combine_time_decision_3():
    df = combine_position()[['交易日期', '指数净值', 'RSJ_5D持仓状态', 'cov_+ret_pct_5D持仓状态', 'covprice_mid持仓状态',
                             'IVdelta_mid_5D持仓状态', 'YTMdelta_mid持仓状态', 'IVnan_pct持仓状态', '等权仓位占比']]
    df['日收益率'] = df['指数净值'] / df['指数净值'].shift(1) - 1
    df.loc[0, '日收益率'] = 0
    df['组合累积净值'] = np.nan
    df.loc[0, '组合累积净值'] = 1
    for i in range(1, len(df)):
        df.loc[i, '组合累积净值'] = df.loc[i - 1, '组合累积净值'] * (1 + df.loc[i, '日收益率'] * df.loc[i, '等权仓位占比'])
    return df

# 组合择时设定4
def combine_time_decision_4():
    df = combine_position()[['交易日期', '指数净值', 'RSJ_5D持仓状态', 'cov_+ret_pct_5D持仓状态', 'covprice_mid持仓状态',
                             'IVdelta_mid_5D持仓状态', 'YTMdelta_mid持仓状态', 'IVnan_pct持仓状态', '分权仓位占比']]
    df['日收益率'] = df['指数净值'] / df['指数净值'].shift(1) - 1
    df.loc[0, '日收益率'] = 0
    df['组合累积净值'] = np.nan
    df.loc[0, '组合累积净值'] = 1
    for i in range(1, len(df)):
        df.loc[i, '组合累积净值'] = df.loc[i - 1, '组合累积净值'] * (1 + df.loc[i, '日收益率'] * df.loc[i, '分权仓位占比'])
    return df

def net_value_drawing(df, name = '组合1 等权 4个及以上满仓 否则空仓', filename = '择时信号组合使用图'):

    matplotlib.rcParams['font.sans-serif'] = ['SimHei']
    matplotlib.rcParams['axes.unicode_minus'] = False

    fig, ax1 = plt.subplots(figsize=(15, 8))

    ax1.plot(df['交易日期'], df['指数净值'], linestyle='--', color='black', label='基准指数')
    selected_ticks = np.arange(0, len(df), 66)
    plt.xticks(selected_ticks, rotation=45, ha='right')
    plt.xlim(df['交易日期'].min(), df['交易日期'].max())
    ax1.plot(df['交易日期'], df.iloc[:, -1], color='red', label='组合累积净值')
    ax1.set_ylabel('净值')

    ax2 = ax1.twinx()
    # ax2.bar(df['交易日期'], df.iloc[:, -3], color='grey', label='组合持仓状态')
    ax2.fill_between(df['交易日期'], df.iloc[:, -3], color='grey', label='组合持仓状态', interpolate=True, alpha=0.5)
    ax2.set_ylim(0, 1.2)

    handles1, labels1 = ax1.get_legend_handles_labels()
    handles2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(handles1 + handles2, labels1 + labels2, loc='upper left')

    plt.title(f'{name}')
    plt.xlabel('时间')
    plt.savefig(f'{filename}\\{name}.jpg')
    plt.close()

def picture_massive_drawing(df1, df2, df3, df4):
    filename = '择时信号组合使用图'
    try:
        os.mkdir(filename)
        print(f'文件夹{filename}创建成功')
    except:
        print(f'{filename}文件夹已存在，无需创建')

    net_value_drawing(df1, '组合1 等权 4个及以上满仓 否则空仓', filename)
    net_value_drawing(df2, '组合2 等权 4个及以上满仓 2个及以下空仓', filename)
    net_value_drawing(df3, '组合3 等权 灵活仓位', filename)
    net_value_drawing(df4, '组合4 不等权 灵活仓位', filename)

if __name__ == '__main__':

    # 获取组合择时信号
    filename = '择时信号组合回测结果'
    try:
        os.mkdir(filename)
        print(f'文件夹{filename}创建成功')
    except:
        print(f'{filename}文件夹已存在，无需创建')
    df1 = combine_time_decision_1()
    df2 = combine_time_decision_2()
    df3 = combine_time_decision_3()
    df4 = combine_time_decision_4()

    df1.to_excel(f'{filename}\\组合1 等权 4个及以上满仓 否则空仓.xlsx', index=False)
    df2.to_excel(f'{filename}\\组合2 等权 4个及以上满仓 2个及以下空仓.xlsx', index=False)
    df3.to_excel(f'{filename}\\组合3 等权 灵活仓位.xlsx', index=False)
    df4.to_excel(f'{filename}\\组合4 不等权 灵活仓位.xlsx', index=False)

    # 获取组合择时信号净值与仓位图
    picture_massive_drawing(df1, df2, df3, df4)

    # 获取择时净值业绩表现指标
    ls = ['组合1 等权 4个及以上满仓 否则空仓', '组合2 等权 4个及以上满仓 2个及以下空仓', '组合3 等权 灵活仓位', '组合4 不等权 灵活仓位']
    df_all = pd.DataFrame()
    for type in ls:
        df = pd.read_excel(f'{filename}\\{type}.xlsx').rename(columns={'组合持仓状态':'持仓状态', '组合累积净值':'累积净值'})
        df = indicator_merge(df)
        df = df.T
        df.columns = df.iloc[0, :]
        df = df.iloc[1:, :]
        df.index = [f'{type}']
        if type in ['组合3 等权 灵活仓位', '组合4 不等权 灵活仓位']:
            df['持有期占比'] = np.nan
            df['进场次数'] = np.nan
            df['看多择时胜率'] = np.nan
            df['多空总择时胜率'] = np.nan
        df_all = pd.concat([df_all, df])

    filename = '组合择时信号评价指标'
    try:
        os.mkdir(filename)
        print(f'文件夹{filename}创建成功')
    except:
        print(f'{filename}文件夹已存在，无需创建')
    df_all.to_excel(f'{filename}/各组合择时信号评价指标.xlsx')
    
    
